import base64
import os

try:
    import rsa
except ImportError:
    os.system('pip install rsa')
    import rsa


class Encryptor(object):

    def encrypt(self, message):
        pass

    def decrypt(self, message):
        pass


class RASEncryptor(Encryptor):
    DEFAULT_PUBLIC_PEM = 'rsa_pub.pem'
    DEFAULT_PRIVATE_PEM = 'rsa_private.pem'

    def __init__(self, public_pem_path=None, private_pem_path=None):
        self.public_key = None
        self.private_key = None
        file_path = os.path.abspath(__file__)
        root_path, _ = os.path.split(file_path)
        public_pem = public_pem_path or os.path.join(root_path, self.DEFAULT_PUBLIC_PEM)
        private_pem = private_pem_path or os.path.join(root_path, self.DEFAULT_PRIVATE_PEM)

        if not os.path.exists(public_pem):
            raise ValueError(f'Encrypt pub pem file not find in {public_pem}')
        if not os.path.exists(private_pem):
            raise ValueError(f'Encrypt private pem file not find in {private_pem}')

        with open(public_pem, mode='rb') as f:
            self.public_Key = rsa.PublicKey.load_pkcs1(f.read())
        with open(private_pem, mode='rb') as f:
            self.private_key = rsa.PrivateKey.load_pkcs1(f.read())

    def encrypt(self, message):
        if message is None or len(message) == 0:
            return message
        encoding = message.encode('utf-8')
        results = b''
        for i in range(0, len(encoding), 20):
            result = rsa.encrypt(encoding[i:i + 20], self.public_Key)
            results += result
        return base64.b64encode(results).decode('utf-8')

    def decrypt(self, message):
        if message is None or len(message) == 0:
            return message
        encoding = base64.b64decode(message.encode('utf-8'))
        results = b''
        for i in range(0, len(encoding), 32):
            result = rsa.decrypt(encoding[i:i + 32], self.private_key)
            results += result
        return results.decode('utf-8')


class EncryptUtil(object):
    SUPPORT_ENCRYPT = {'rsa_cloud_test': {'flag': '{cipher}{rsa_cloud_test}', 'encryptor': RASEncryptor()}}
    END_FLAG = '{END}'
    START_FLAG = '{cipher}'

    @classmethod
    def encrypt_string(cls, message, encrypt_type=None):
        # 确保为字符串
        message = str(message)
        encrypt_type = encrypt_type or 'rsa_cloud_test'
        if encrypt_type not in cls.SUPPORT_ENCRYPT.keys():
            raise ValueError(f'Not support encrypt type {encrypt_type}')

        encrypt_dict = cls.SUPPORT_ENCRYPT[encrypt_type]
        flag = encrypt_dict['flag']
        encryptor = encrypt_dict['encryptor']

        encrypt_string = encryptor.encrypt(message)

        return flag + encrypt_string + cls.END_FLAG

    @classmethod
    def decrypt_string(cls, message):
        # 确保为字符串
        message = str(message)

        result = ''
        encrpyt_string, before, message = cls._get_encrpyt_string(message)
        while encrpyt_string:
            result += before + cls._decrypt(encrpyt_string)
            encrpyt_string, before, message = cls._get_encrpyt_string(message)
        result += message

        return result

    @classmethod
    def _decrypt(cls, encrypt_string):
        for encrypt_dict in cls.SUPPORT_ENCRYPT.values():
            flag = encrypt_dict['flag']
            if encrypt_string.startswith(flag):
                encryptor = encrypt_dict['encryptor']
                result = encryptor.decrypt(encrypt_string[len(flag):])
                return result
        # 如果没有合适的解密方法，直接返回原内容
        return encrypt_string

    @classmethod
    def _get_encrpyt_string(cls, message):
        # 返回加密字符串，前部字符串，后部字符串
        start_index = message.find(cls.START_FLAG, 0)
        if start_index >= 0:
            end_index = message.find(cls.END_FLAG, start_index)
            if end_index >= 0:
                return message[start_index:end_index], message[0:start_index], message[end_index + len(cls.END_FLAG):]
            else:
                return message[start_index:], message[0:start_index], ''
        else:
            return None, '', message


if __name__ == '__main__':
    # pubkey, privateKey = rsa.newkeys(256)  # 生成随机秘钥
    # with open('rsa_pub.pem', mode='wb') as f:
    #     f.write(pubkey.save_pkcs1())
    # with open('rsa_private.pem', mode='wb') as f:
    #     f.write(privateKey.save_pkcs1())

    result = EncryptUtil.encrypt_string("我的hello word!")
    print(result)
    code = "{rsa_cloest}PzzTaUOn8q2xhmCG5cBuhtw4UPfcv8bexa/TdCldBgU={END}"
    print(EncryptUtil.decrypt_string(code))
    code2 = "{cipher}{rsa_cloud_test}agXhSs0LD1vURreyh8Mv5717FExIFekxacX4CZqQbAY={END}"
    print(EncryptUtil.decrypt_string(code2))


